// main.go --------------------------------------------------
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"github.com/sbinet/pstree"
)

func main() {
	log.SetFlags(0)
	
	pid := flag.Int("pid", os.Getpid(), "目标进程PID")
	jsonOutput := flag.Bool("json", false, "启用JSON输出格式")
	flag.Parse()

	tree, err := pstree.BuildTree()
	if err != nil {
		log.Printf("警告: 部分进程解析失败 (%v)", err)
	}

	proc, exists := tree.Procs[*pid]
	if !exists {
		log.Fatalf("错误: PID %d 不存在", *pid)
	}

	if *jsonOutput {
		outputJSON(tree, proc)
	} else {
		outputText(tree, proc)
	}
}

func outputJSON(tree *pstree.ProcessTree, proc *pstree.Process) {
	result := map[string]interface{}{
		"pid":        proc.PID,
		"name":       proc.Name,
		"cmdline":    strings.TrimSpace(proc.CmdLine),
		"parent_pid": proc.PPID,
		"children":   getChildren(tree, proc.PID),
		"process_tree": tree.FormatChain(proc.PID),
	}

	enc := json.NewEncoder(os.Stdout)
	enc.SetIndent("", "  ")
	if err := enc.Encode(result); err != nil {
		log.Fatal("JSON编码失败:", err)
	}
}

func outputText(tree *pstree.ProcessTree, proc *pstree.Process) {
	fmt.Printf("进程树 %d\n", proc.PID)
	fmt.Printf("名称:    %s\n", proc.Name)
	fmt.Printf("命令行:  %s\n", strings.TrimSpace(proc.CmdLine))
	fmt.Printf("父PID:   %d\n", proc.PPID)
	fmt.Println("子进程结构:")
	displayTree(proc.PID, tree, 1)
}

// 树形显示辅助函数 -------------------------------------------
func displayTree(pid int, tree *pstree.ProcessTree, depth int) {
	proc := tree.Procs[pid]
	indent := strings.Repeat("│   ", depth-1)

	for i, cid := range proc.Children {
		lastChild := i == len(proc.Children)-1
		if lastChild {
			fmt.Printf("%s└── %d(%s)\n", indent, cid, tree.Procs[cid].Name)
		} else {
			fmt.Printf("%s├── %d(%s)\n", indent, cid, tree.Procs[cid].Name)
		}
		displayTree(cid, tree, depth+1)
	}
}

func getChildren(tree *pstree.ProcessTree, pid int) []int {
	proc, exists := tree.Procs[pid]
	if !exists {
		return nil
	}
	return proc.Children
}
